import os
import torch

from .plugins_base import PluginBase
from trainer.trainer import PluginType, TrainContext


class PluginLoadModel(PluginBase):
    plugin_hooks = {
        PluginType.TRAIN_BEGIN: "auto_load"
    }

    def __init__(self, checkpoint_path: str = None):
        self.checkpoint_path = checkpoint_path

    def auto_load(self, ctx: TrainContext):
        results_dir = "./results"

        if self.checkpoint_path and os.path.isfile(self.checkpoint_path):
            latest_ckpt = self.checkpoint_path
            latest_dir = os.path.dirname(os.path.dirname(latest_ckpt))
        else:
            if not os.path.isdir(results_dir):
                raise FileNotFoundError(f"[AutoResumePlugin] No results folder found at: {results_dir}")

            subdirs = [
                os.path.join(results_dir, d)
                for d in os.listdir(results_dir)
                if os.path.isdir(os.path.join(results_dir, d)) and "_" in d
            ]
            if not subdirs:
                raise FileNotFoundError("[AutoResumePlugin] No valid training folders found in ./results")

            latest_dir = max(subdirs, key=os.path.getmtime)
            checkpoints_dir = os.path.join(latest_dir, "checkpoints")

            if not os.path.isdir(checkpoints_dir):
                raise FileNotFoundError(f"[AutoResumePlugin] No checkpoints folder in: {latest_dir}")

            ckpt_files = [
                os.path.join(checkpoints_dir, f)
                for f in os.listdir(checkpoints_dir)
                if f.endswith(".pth")
            ]
            if not ckpt_files:
                raise FileNotFoundError(f"[AutoResumePlugin] No checkpoint files found in: {checkpoints_dir}")

            latest_ckpt = max(ckpt_files, key=os.path.getmtime)
            latest_dir = os.path.dirname(os.path.dirname(latest_ckpt))

        checkpoint = torch.load(latest_ckpt, map_location=ctx.device)

        ctx.model.load_state_dict(checkpoint["model_state_dict"])
        ctx.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        ctx.scheduler.load_state_dict(checkpoint["scheduler_state_dict"])

        start_epoch = checkpoint["epoch"] + 1
        best_miou = checkpoint.get("miou", 0.0)

        folder_name = os.path.basename(latest_dir)
        if "_" in folder_name:
            timestamp, short_uuid = folder_name.split("_")
            ctx.workspace["uuid"] = short_uuid
            ctx.workspace["start_time"] = timestamp
        
        ctx.workspace["results_folder"] = latest_dir
        ctx.workspace["start_epoch"] = start_epoch
        ctx.workspace["best_miou"] = best_miou
